import argparse
import os
import mimetypes
import warnings
import langdetect

# Setup environment variables early
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['TRANSFORMERS_CACHE'] = './models_cache'
warnings.filterwarnings("ignore")

import whisper
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer, AutoModelForImageClassification, AutoImageProcessor

class AudioEmotionRecognizer:
    def __init__(self):
        print("初始化音频情感识别器...")
        self.asr_model = whisper.load_model("base")

    def audio_to_text(self, audio_path):
        result = self.asr_model.transcribe(audio_path)
        return result["text"]
    
    def process_file(self, audio_path):
        if not os.path.exists(audio_path):
            return {"error": f"文件不存在: {audio_path}"}
            
        print(f"正在识别音频中的语音: {audio_path}")
        text = self.audio_to_text(audio_path)
        print(f"识别文本: '{text}'")
        
        if not text.strip():
            return {"error": "无法识别音频中的语音或语音内容为空"}
        
        detected_lang = 'unknown'
        if text.strip():
            try:
                detected_lang = langdetect.detect(text)
            except Exception as e:
                print(f"警告: 无法检测音频转录文本语言，将尝试使用英文情感模型。错误: {e}")
                detected_lang = 'en'
        else:
            detected_lang = 'en'

        print(f"检测到转录文本语言: {detected_lang}")

        text_emotion_recognizer = None
        if detected_lang == 'zh-cn':
            text_emotion_recognizer = ChineseTextEmotionRecognizer()
        else:
            text_emotion_recognizer = EnglishTextEmotionRecognizer()

        print("正在分析文本情感...")
        
        raw_emotion_result = text_emotion_recognizer.analyze_emotion(text)
        
        emotion_label = None
        emotion_score = None

        if isinstance(raw_emotion_result, dict):
            if 'label' in raw_emotion_result and 'score' in raw_emotion_result:
                emotion_label = raw_emotion_result['label']
                emotion_score = float(raw_emotion_result['score'])
            elif 'emotion_label' in raw_emotion_result and 'emotion_score' in raw_emotion_result:
                emotion_label = raw_emotion_result['emotion_label']
                emotion_score = float(raw_emotion_result['emotion_score'])
        
        if emotion_label is None:
            return {"error": "情感分析结果格式不正确。"}

        return {
            "type": "audio",
            "text": text,
            "emotion_label": emotion_label,
            "emotion_score": emotion_score
        }

from modelscope.pipelines import pipeline as ms_pipeline
from modelscope.utils.constant import Tasks as MSTasks

class ChineseTextEmotionRecognizer:
    def __init__(self):
        print("初始化中文文本情感识别器...")
        self.classifier = ms_pipeline(MSTasks.text_classification, 'iic/nlp_structbert_emotion-classification_chinese-base', model_revision='v1.0.0')
        print("中文文本情感模型加载成功 (ModelScope)")

    def analyze_emotion(self, text):
        input_data = {"text": text}
        result = self.classifier(input_data)
        
        highest_score_index = result['scores'].index(max(result['scores']))
        highest_label = result['labels'][highest_score_index]
        highest_score = result['scores'][highest_score_index]
        
        return {"emotion_label": highest_label, "emotion_score": highest_score}

    def process_file(self, file_path):
        if not os.path.exists(file_path):
            return {"error": f"文件不存在: {file_path}"}

        with open(file_path, 'r', encoding='utf-8') as file:
            text_content = file.read().strip()
        
        if not text_content.strip():
            return {"error": "文本内容为空"}
            
        print(f"正在分析中文文本情感...")
        result = self.analyze_emotion(text_content)
        return {
            "type": "text_chinese",
            "text": text_content,
            "emotion_label": result['emotion_label'],
            "emotion_score": float(result['emotion_score'])
        }

class EnglishTextEmotionRecognizer:
    def __init__(self):
        print("初始化英文文本情感识别器...")
        model_name = "distilbert-base-uncased-emotion"
        local_model_dir = f"./models/{model_name.replace('/', '_')}"

        try:
            tokenizer = AutoTokenizer.from_pretrained(local_model_dir, local_files_only=True)
            model = AutoModelForSequenceClassification.from_pretrained(local_model_dir, local_files_only=True)
            print(f"英文文本情感模型从本地加载成功: {local_model_dir}")
        except Exception as e:
            print(f"英文文本情感模型本地加载失败: {e}. 尝试在线下载并保存...")
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForSequenceClassification.from_pretrained(model_name)
            model.save_pretrained(local_model_dir)
            tokenizer.save_pretrained(local_model_dir)
            print(f"英文文本情感模型已下载并保存到: {local_model_dir}")

        self.emotion_classifier = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer
        )

    def analyze_emotion(self, text):
        return self.emotion_classifier(text)[0]

    def process_file(self, file_path):
        if not os.path.exists(file_path):
            return {"error": f"文件不存在: {file_path}"}

        with open(file_path, 'r', encoding='utf-8') as file:
            text_content = file.read().strip()
        
        if not text_content.strip():
            return {"error": "文本内容为空"}
            
        print(f"正在分析英文文本情感...")
        result = self.analyze_emotion(text_content)
        return {
            "type": "text_english",
            "text": text_content,
            "emotion_label": result['label'],
            "emotion_score": float(result['score'])
        }

import PIL.Image
class ImageEmotionRecognizer:
    def __init__(self):
        print("初始化图片情感识别器...")
        model_name = "dima806/facial_emotions_image_detection"
        local_model_dir = f"./models/{model_name.replace('/', '_')}"

        try:
            processor = AutoImageProcessor.from_pretrained(local_model_dir, local_files_only=True)
            model = AutoModelForImageClassification.from_pretrained(local_model_dir, local_files_only=True)
            print(f"图片情感模型从本地加载成功: {local_model_dir}")
        except Exception as e:
            print(f"图片情感模型本地加载失败: {e}. 尝试在线下载并保存...")
            processor = AutoImageProcessor.from_pretrained(model_name)
            model = AutoModelForImageClassification.from_pretrained(model_name)
            model.save_pretrained(local_model_dir)
            processor.save_pretrained(local_model_dir)
            print(f"图片情感模型已下载并保存到: {local_model_dir}")

        self.emotion_classifier = pipeline(
            "image-classification",
            model=model,
            feature_extractor=processor
        )

    def process_file(self, image_path):
        if not os.path.exists(image_path):
            return {"error": f"文件不存在: {image_path}"}
        
        print(f"正在分析图片情感: {image_path}")
        image = PIL.Image.open(image_path).convert("RGB")
        result = self.emotion_classifier(image)[0]
        return {
            "type": "image",
            "emotion_label": result['label'],
            "emotion_score": float(result['score'])
        }

import cv2
from collections import Counter
class VideoEmotionRecognizer:
    def __init__(self, interval=10):
        print("初始化视频情感识别器...")
        self.interval = interval
        
        model_name = "dima806/facial_emotions_image_detection"
        local_model_dir = f"./models/{model_name.replace('/', '_')}"

        try:
            processor = AutoImageProcessor.from_pretrained(local_model_dir, local_files_only=True)
            model = AutoModelForImageClassification.from_pretrained(local_model_dir, local_files_only=True)
            print(f"视频情感模型 (内部图片模型) 从本地加载成功: {local_model_dir}")
        except Exception as e:
            print(f"视频情感模型 (内部图片模型) 本地加载失败: {e}. 尝试在线下载并保存...")
            processor = AutoImageProcessor.from_pretrained(model_name)
            model = AutoModelForImageClassification.from_pretrained(model_name)
            model.save_pretrained(local_model_dir)
            processor.save_pretrained(local_model_dir)
            print(f"视频情感模型 (内部图片模型) 已下载并保存到: {local_model_dir}")

        self.emotion_classifier = pipeline(
            "image-classification",
            model=model,
            feature_extractor=processor
        )

    def process_file(self, video_path):
        if not os.path.exists(video_path):
            return {"error": f"文件不存在: {video_path}"}

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            return {"error": f"无法打开视频文件: {video_path}"}

        frame_count = 0
        all_results = []  # 存储 (label, score) 元组
        
        print(f"开始处理视频: {video_path}, 每 {self.interval} 帧抽取一帧...")
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            if frame_count % self.interval == 0:
                img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_img = PIL.Image.fromarray(img)
                try:
                    pred = self.emotion_classifier(pil_img)
                    label = pred[0]['label']
                    score = pred[0]['score']
                    all_results.append((label, score))
                except Exception as e:
                    print(f"警告: 处理帧 {frame_count} 时出错: {e}")
            frame_count += 1

        cap.release()

        if not all_results:
            return {"error": "未识别到任何有效帧或情感。"}

        # 计算每种情感的统计信息
        emotion_stats = {}
        for label, score in all_results:
            if label not in emotion_stats:
                emotion_stats[label] = {'count': 0, 'scores': []}
            emotion_stats[label]['count'] += 1
            emotion_stats[label]['scores'].append(score)

        # 找出最常见的情感
        most_common_emotion = max(emotion_stats.keys(), key=lambda x: emotion_stats[x]['count'])
        
        # 计算主要情感的平均置信度
        main_emotion_avg_score = sum(emotion_stats[most_common_emotion]['scores']) / len(emotion_stats[most_common_emotion]['scores'])
        
        # 构建情感分布（包含平均置信度）
        emotion_distribution = {}
        for label, stats in emotion_stats.items():
            avg_score = sum(stats['scores']) / len(stats['scores'])
            emotion_distribution[label] = {
                'count': stats['count'],
                'avg_confidence': avg_score
            }

        return {
            "type": "video",
            "main_emotion_label": most_common_emotion,
            "main_emotion_confidence": main_emotion_avg_score,
            "total_frames_analyzed": len(all_results),
            "emotion_distribution": emotion_distribution
        }

def detect_file_type(file_path):
    mime_type, _ = mimetypes.guess_type(file_path)
    
    if mime_type:
        if mime_type.startswith('text/'):
            return "text"
        elif mime_type.startswith('image/'):
            return "image"
        elif mime_type.startswith('video/'):
            return "video"
        elif mime_type.startswith('audio/'):
            return "audio"
    
    ext = os.path.splitext(file_path)[1].lower()
    if ext in ['.txt']:
        return "text"
    elif ext in ['.png', '.jpg', '.jpeg', '.bmp', '.gif']:
        return "image"
    elif ext in ['.mp4', '.avi', '.mov', '.mkv']:
        return "video"
    elif ext in ['.wav', '.mp3', '.flac', '.m4a']:
        return "audio"
        
    return "unknown"

def main():
    parser = argparse.ArgumentParser(description='统一情感识别系统')
    parser.add_argument('file_path', type=str, help='要处理的文件路径')
    parser.add_argument('--video_interval', type=int, default=10,
                        help='如果文件是视频，每隔多少帧抽取一帧进行识别（默认10）。')
    args = parser.parse_args()

    file_path = args.file_path
    
    if not os.path.exists(file_path):
        print(f"错误: 文件不存在 '{file_path}'")
        return

    file_type = detect_file_type(file_path)
    print("\n--- 文件检测 ---")
    print(f"检测到文件类型: {file_type}")

    result = None
    if file_type == "text":
        text_sample = ""
        with open(file_path, 'r', encoding='utf-8') as f:
            text_sample = f.read(1000)
        
        detected_lang = 'unknown'
        if text_sample.strip():
            try:
                detected_lang = langdetect.detect(text_sample)
            except Exception as e:
                print(f"警告: 无法检测文本语言，将尝试使用英文模型。错误: {e}")
                detected_lang = 'en'

        print(f"检测到文本语言: {detected_lang}")

        if detected_lang == 'zh-cn':
            recognizer = ChineseTextEmotionRecognizer()
        else:
            recognizer = EnglishTextEmotionRecognizer()
        result = recognizer.process_file(file_path)
    elif file_type == "image":
        recognizer = ImageEmotionRecognizer()
        result = recognizer.process_file(file_path)
    elif file_type == "video":
        recognizer = VideoEmotionRecognizer(interval=args.video_interval)
        result = recognizer.process_file(file_path)
    elif file_type == "audio":
        recognizer = AudioEmotionRecognizer()
        result = recognizer.process_file(file_path)
    else:
        print(f"不支持的文件类型: {file_type}")
        return

    if result and "error" in result:
        print(f"\n处理文件时发生错误: {result['error']}")
    elif result:
        print("\n--- 情感识别结果 ---")
        if result['type'] == 'text_english' or result['type'] == 'text_chinese':
            print(f"类型: {result['type'].replace('text_', '')} 文本")
            print(f"原文: '{result['text'][:100]}...'")
            print(f"情感: {result['emotion_label']}")
            print(f"置信度: {result['emotion_score']:.2%}")
        elif result['type'] == 'image':
            print(f"类型: 图片")
            print(f"情感: {result['emotion_label']}")
            print(f"置信度: {result['emotion_score']:.2%}")
        elif result['type'] == 'video':
            print(f"类型: 视频")
            print(f"主要情感: {result['main_emotion_label']}")
            print(f"主要情感置信度: {result['main_emotion_confidence']:.2%}")
            print(f"分析帧数: {result['total_frames_analyzed']}")
            print("情感分布:")
            for label, stats in result['emotion_distribution'].items():
                print(f"  {label}: {stats['count']} 帧 (平均置信度: {stats['avg_confidence']:.2%})")
        elif result['type'] == 'audio':
            print(f"类型: 音频")
            print(f"识别文本: '{result['text'][:100]}...' ")
            print(f"情感: {result['emotion_label']}")
            print(f"置信度: {result['emotion_score']:.2%}")
    else:
        print("未获取到结果。")

if __name__ == "__main__":
    # 首次运行请确保已安装所有必要的Python库
    print("\n--- 首次运行注意 ---")
    print("请确保已安装以下Python库:")
    print("pip install transformers torch torchaudio opencv-python Pillow modelscope openai-whisper langdetect")
    print("-------------------\n")
    main() 